package utils

import chisel3._
import chisel3.util._
import chisel3.util.experimental._
import chiseltest._

import chiseltest.formal._
import chiseltest.experimental._
import chisel3.stage.{ChiselStage, ChiselGeneratorAnnotation}

import org.scalatest.Tag
import org.scalatest.flatspec.AnyFlatSpec

import chiseltest.ChiselScalatestTester
import chiseltest.formal.backends.FormalEngineAnnotation

class FormalSRAMSpec extends AnyFlatSpec with ChiselScalatestTester with Formal {

  "SRAM" should "pass" in {
    verify(new SRAMSpec, Seq(BoundedCheck(50), BtormcEngineAnnotation))
  }

  "SRAM w/ arbiter" should "pass" in {
    verify(new SRAMArbSpec, Seq(BoundedCheck(50), BtormcEngineAnnotation))
  }

}

class SRAMSpec() extends Module {


  val ways = 2
  val sets = 16
  val DataType = UInt(8.W)
  val WordType = Vec(ways, DataType)

  val sram = Module(new SRAMTemplate(DataType, sets, ways, false, false, false));
  val port = IO(new Bundle {
    val r = Flipped(new SRAMReadBus(DataType, sets, ways))
    val w = Flipped(new SRAMWriteBus(DataType, sets, ways))
  })

  sram.io <> port

  val saveVal  = RegInit(Bool(), false.B);
  val saveAddr = RegInit(chiselTypeOf(port.w.req.bits.setIdx), 0.U);
  val saveData = RegInit(DataType, 0.U);
  val saveMask = RegInit(UInt(ways.W), 0.U);

  saveVal  := port.w.req.valid;
  saveAddr := port.w.req.bits.setIdx;
  saveData := port.w.req.bits.data;


  // only works when SyncReadMem is WriteFirst, for read & write on same cycle.
  // when(past(port.w.req.valid && port.r.req.valid && port.r.req.bits.addr === port.w.req.bits.addr)) {
  //   assert(port.r.resp.data === past(port.w.req.bits.data));
  // }


  val popCountMask = PopCount(port.w.req.bits.waymask.getOrElse(1.U));
  val priorEncMask = PriorityEncoder(port.w.req.bits.waymask.getOrElse(1.U));

  // the delay between req and res of a read 
  // assert(false.B)
  // assert(!past(port.r.req.valid && !port.w.req.valid))
  when (past(port.r.req.valid && !port.w.req.valid)) {
    assert(port.r.resp.addr === past(port.r.req.bits.setIdx))
  }

  // successive write & read request to check if data is stored in specific addr
  when (
    past(port.w.req.valid && !port.r.req.valid, 2) &&
    past(!port.w.req.valid && port.r.req.valid) &&
    past(port.w.req.bits.setIdx, 2) === past(port.r.req.bits.setIdx) &&
    past(popCountMask, 2) === 1.U   
  ) {
    assert(port.r.resp.data(past(priorEncMask, 2)).asUInt === past(port.w.req.bits.data.asUInt, 2));
  }


  when(
    saveVal && 
    past(!port.w.req.valid && port.r.req.valid) && 
    saveAddr === past(port.r.req.bits.setIdx) &&
    past(popCountMask, 2) === 1.U
  ) {
    assert(port.r.resp.data(past(priorEncMask, 2)).asUInt === saveData);
  }
}

class SRAMArbSpec extends Module {

  val ways = 2
  val sets = 16
  val DataType = UInt(8.W)
  val WordType = Vec(ways, DataType)
  val nRead = 2

  val sram = Module(new SRAMTemplateWithArbiter(nRead, DataType, sets, ways, false));
  val port = IO(new Bundle {
    val r0 = Flipped(new SRAMReadBus(DataType, sets, ways))
    val r1 = Flipped(new SRAMReadBus(DataType, sets, ways))
    val w = Flipped(new SRAMWriteBus(DataType, sets, ways))
  })

  sram.io.r(0) <> port.r0
  sram.io.r(1) <> port.r1
  sram.io.w <> port.w

  val popCountMask = PopCount(port.w.req.bits.waymask.getOrElse(1.U));
  val priorEncMask = PriorityEncoder(port.w.req.bits.waymask.getOrElse(1.U));

  // the delay between req and res of a read 
  // assert(past(port.r.req.valid && !port.w.req.valid))
 
  when (
    past(port.w.req.valid && !port.r0.req.valid, 2) &&
    past(!port.w.req.valid && port.r0.req.valid) &&
    past(port.w.req.bits.setIdx, 2) === past(port.r0.req.bits.setIdx) &&
    past(popCountMask, 2) === 1.U   
  ) {
    assert(port.r0.resp.data(past(priorEncMask, 2)).asUInt === past(port.w.req.bits.data.asUInt, 2));
  }

}
